// Copyright 2009-2024 NTESS. Under the terms
// of Contract DE-NA0003525 with NTESS, the U.S.
// Government retains certain rights in this software.
//
// Copyright (c) 2009-2024, NTESS
// All rights reserved.
//
// Portions are copyright of other developers:
// See the file CONTRIBUTORS.TXT in the top level directory
// of the distribution for more information.
//
// This file is part of the SST software package. For license
// information, see the LICENSE file in the top level directory of the
// distribution.


#include <sst_config.h>
#include <sst/core/link.h>
#include "sst/elements/memHierarchy/util.h"
#include "membackend/accurateDRAMBackend.h"
#include "membackend/accurateAddressMappingConfig.h"

using namespace SST;
using namespace SST::MemHierarchy;

/*------------------------------- Simple DRAM ------------------------------- */
/* SimpleDRAM is a very simplified DRAM timing model. It considers only the
 * latency differences between accessing open/closed rows and implements banking.
 * The latencies are as follows:
 *      Correct row already open: tCAS
 *      No row open: tRCD + tCAS
 *      Wrong row open: tRP + tRCD + tCAS
 *
 *  SimpleDRAM maps address to banks and rows as follows
 *      Example: 64B lines, 8 banks, 8KiB rows (256 columns/row)
 *      |...  19|18     9|8    6|5    0|
 *      |  Row  | Column | Bank | Line |
 *
 * Implementation notes:
 *  SimpleDRAM relies on external (memController or other backend) request limiting. SimpleDRAM will block banks but will accept multiple requests per cycle if they don't have a bank conflict.
 *  SimpleDRAM assumes transfer granularity is a cache line -> this can be changed by changing 'cache_line_size_in_bytes'
 *  SimpleDRAM does not model bus contention back to the controller -> possible for multiple requests to return in the same cycle
 *
 *  TODO
 *   * Randomize bank bits?
 *   * Decide whether to consider tRAS
 *   * Allow multiple ranks?
 */


AccurateDRAM::AccurateDRAM(ComponentId_t id, Params &params) : SimpleMemBackend(id, params){ 
    // Get parameters
    m_cycle = 0;
    tCAS = params.find<uint64_t>("tCAS", 9);
    tRCD = params.find<uint64_t>("tRCD", 9);
    tRP = params.find<uint64_t>("tRP", 9);
    tRAS = params.find<uint64_t>("tRAS", 10); // tRAS defaults to 0 if not specified

    // 新增参数
    tAL = params.find<uint64_t>("tAL", 0);
    tCL = params.find<uint64_t>("tCL", 9); // tCL defaults to tCAS if not specified
    tCWL = params.find<uint64_t>("tCWL", tCL); // tCWL defaults to tCL if not specified
    tBurstCycle = params.find<uint64_t>("burst_length", 8) / 2; // Default burst length is 8, divided by 2 for cycles
    std::cout << "tAL: " << tAL << ", tCL: " << tCL << ", tCWL: " << tCWL << ", tBurstCycle: " << tBurstCycle << std::endl;

    // 协议信息
    std::string protocol = params.find<std::string>("protocol", "DDR4");

    // DDR3专用参数
    uint64_t tRRD = params.find<uint64_t>("tRRD", 4);

    // DDR4专用参数
    uint64_t tRRD_L = params.find<uint64_t>("tRRD_L", 4);
    uint64_t tRRD_S = params.find<uint64_t>("tRRD_S", 2);

    uint64_t tFAW = params.find<uint64_t>("tFAW", 16); // Fast access window, if applicable
    
    std::string cycTime = params.find<std::string>("cycle_time", "4ns");
    uint64_t banks = params.find<uint64_t>("banks", 8);
    UnitAlgebra lineSize(params.find<std::string>("bank_interleave_granularity", "64B"));
    UnitAlgebra rowSize(params.find<std::string>("row_size", "8KiB"));

    uint64_t bankgroups = params.find<uint64_t>("bankgroups", 1);
    uint64_t bank_per_group = params.find<uint64_t>("bank_per_group", 1);
    UnitAlgebra rows_size(params.find<std::string>("rows", "4KiB"));
    uint64_t rows = params.find<uint64_t>("rows", 4);
    uint64_t columns = params.find<uint64_t>("columns", 256);
    uint64_t bus_width = params.find<uint64_t>("bus_width", 64);
    uint64_t device_width = params.find<uint64_t>("device_width", 8);

    uint64_t trans_queue_size = params.find<uint64_t>("trans_queue_size", 32);
    uint64_t cmd_queue_size = params.find<uint64_t>("cmd_queue_size", 8);

    string my_row_policy = params.find<std::string>("row_policy", "closed");
    string my_queue_policy = params.find<std::string>("queue_policy", "per_bank");
    uint64_t BL = params.find<uint64_t>("BL", 8);

    MyRowPolicy my_rowpolicy = MyRowPolicy::INVALID;
    MyQueuePolicy my_queuepolicy = MyQueuePolicy::INVALID;

    // 属性，传params
    if (my_row_policy == "closed") {
        my_rowpolicy = MyRowPolicy::CLOSED;
    } else if (my_row_policy == "open") {
        my_rowpolicy = MyRowPolicy::OPEN;
    } else {
        output->fatal(CALL_INFO, -1, "Invalid row policy '%s'. Must be 'closed' or 'open'.\n", my_row_policy.c_str());
    }

    if(my_queue_policy == "per_bank") {
        my_queuepolicy = MyQueuePolicy::PER_BANK;
    } else if (my_queue_policy == "per_rank") {
        my_queuepolicy = MyQueuePolicy::PER_RANK;
    } else {
        output->fatal(CALL_INFO, -1, "Invalid queue policy '%s'. Must be 'per_bank' or 'per_rank'.\n", my_queue_policy.c_str());
    }

    PolicyParameters policyParams = {
        my_rowpolicy, my_queuepolicy, trans_queue_size, cmd_queue_size
    };
    
    std::cout << "my_queuepolicy: " << static_cast<int>(policyParams.queuePolicy) << std::endl;



    TimeParameters timeParams = {
        tRCD, tAL, tCL, tCWL, tRP, tRAS, tBurstCycle, tRRD, tRRD_L, tRRD_S
    };

    AddressMappingParameters addressParams = {
        this->m_memSize, bus_width, device_width, bankgroups, bank_per_group, rows, columns, BL
    };

    addressMappingConfig = new AddressMappingConfig(&timeParams, &addressParams, &policyParams);
    std::cout << "In addmap, queue_policy is: " << static_cast<int>(addressMappingConfig->getQueuePolicy()) << std::endl;
    controller = new Controller(addressMappingConfig);
    controller->Init();

    bool found = false;
    std::string policyStr = params.find<std::string>("row_policy", "closed", found);
    int verbose = params.find<int>("verbose", 0);

    output = new Output("AccurateDRAM[@p:@l]: ", verbose, 0, Output::STDOUT);

    // Check parameters
    // Supported policies are 'open', 'closed' or 'dynamic'
    if (policyStr != "closed" && policyStr != "open") {
        output->fatal(CALL_INFO, -1, "Invalid param(%s): row_policy - must be 'closed' or 'open'. You specified '%s'.\n", getName().c_str(), policyStr.c_str());
    }

    if (policyStr == "closed") policy = RowPolicy::CLOSED;
    else policy = RowPolicy::OPEN;

    // banks needs to be a power of 2 -> use to set bank mask
    if (!isPowerOfTwo(banks)) {
        output->fatal(CALL_INFO, -1, "Invalid param(%s): banks - must be a power of two. You specified %" PRIu64 ".\n", getName().c_str(), banks);
    }
    bankMask = banks - 1;

    // line size needs to be a power of 2 and have units of bytes
    if (!(lineSize.hasUnits("B"))) {
        output->fatal(CALL_INFO, -1, "Invalid param(%s): cache_line_size_in_bytes - must have units of 'B' (bytes). The units you specified were '%s'.\n", getName().c_str(), lineSize.toString().c_str());
    }
    if (!isPowerOfTwo(lineSize.getRoundedValue())) {
        output->fatal(CALL_INFO, -1, "Invalid param(%s): cache_line_size_in_bytes - must be a power of two. You specified %s.\n", getName().c_str(), lineSize.toString().c_str());
    }
    lineOffset = log2Of(lineSize.getRoundedValue());

    // row size (# columns) needs to be power of 2 and have units of bytes
    if (!(rowSize.hasUnits("B"))) {
        output->fatal(CALL_INFO, -1, "Invalid param(%s): row_size - must have units of 'B' (bytes). You specified %s.\n", getName().c_str(), rowSize.toString().c_str());
    }
    if (!isPowerOfTwo(rowSize.getRoundedValue())) {
        output->fatal(CALL_INFO, -1, "Invalid param(%s): row_size - must be a power of two. You specified %s.\n", getName().c_str(), rowSize.toString().c_str());
    }
    rowOffset = log2Of(rowSize.getRoundedValue());
    std::cout << "AccurateDRAM rowOffset: " << rowOffset << std::endl;

    // Bookkeeping for bank/row state
    busy = (bool*) malloc(sizeof(bool) * banks);
    openRow = (int*) malloc(sizeof(int) * banks);
    for (int i = 0; i < banks; i++) {
        openRow[i] = -1;
        busy[i] = false;
    }

    // Self link for timing requests
    self_link = configureSelfLink("Self", cycTime, new Event::Handler<AccurateDRAM>(this, &AccurateDRAM::handleSelfEvent));

    // Some statistics
    statRowHit = registerStatistic<uint64_t>("row_already_open");
    statRowMissNoRP = registerStatistic<uint64_t>("no_row_open");
    statRowMissRP = registerStatistic<uint64_t>("wrong_row_open");
}

/*
 *  Return response and free bank
 */
// void AccurateDRAM::handleSelfEvent(SST::Event *event){
//     MemCtrlEvent *ev = static_cast<MemCtrlEvent*>(event);
//     if ( ! ev->close ) {
//         if (policy == RowPolicy::CLOSED) {
//             self_link->send(tRP, new MemCtrlEvent(ev->bank));
//         } else {
//             busy[ev->bank] = false;
//         }
//         handleMemResponse(ev->reqId);
//         delete event;
//     } else {
//         openRow[ev->bank] = -1;
//         busy[ev->bank] = false;
//         delete event;
//     }
// }


void AccurateDRAM::handleSelfEvent(SST::Event *event){

    std::cout << "AccurateDRAM handleSelfEvent called and cycle is:" << m_cycle << std::endl;
    MemCtrlEvent *ev = static_cast<MemCtrlEvent*>(event);
    

    ReqId reqId = ev->reqId;
    
    Transaction transaction = controller->getTransactionFromId(reqId);
    
    uint64_t addr = transaction.getAddress();
    TRANSACTION_TYPE type = transaction.getType();

    
    controller->DeleteTransactionInAll(reqId);
    controller->DeleteTransactionBuffer(reqId, type);
    
    controller->DeletePending(reqId, type, addr);
    controller->DeleleReturn(reqId);
    controller->DeleteCommandFromQueue(transaction.getRankId(), transaction.getBankGroupId(), transaction.getBankId(), reqId);

    std::cout << "AccurateDRAM handleSelfEvent: Deleting transaction with reqId: " << reqId << std::endl;
    handleMemResponse(reqId);
    delete ev;
}


bool AccurateDRAM::issueRequest( ReqId reqId, Addr addr, bool isWrite, unsigned numBytes ){

    // Determine bank & row for address
    //  Basic mapping: interleave cache lines across banks

    // std::cout << "type is: " << std::boolalpha << isWrite << std::endl;

    Transaction temp_trans(reqId, addr, isWrite ? TRANSACTION_TYPE::WRITE : TRANSACTION_TYPE::READ, addressMappingConfig);

    if((!controller->WillAcceptTransaction(isWrite)) || !controller->WillAcceptCommand(temp_trans.getRankId(), temp_trans.getBankGroupId(), temp_trans.getBankId())){
        
        std::cout << controller->WillAcceptTransaction(isWrite) << " " <<  controller->WillAcceptCommand(temp_trans.getRankId(), temp_trans.getBankGroupId(), temp_trans.getBankId()) << std::endl;

        std::cout << "## AccurateDRAM issueRequest: Will not accept transaction for address 0x" 
                  << std::setfill('0') << std::setw(8) 
                  << std::hex << addr << std::endl;
        return false;
    }else{
        std::cout << "## AccurateDRAM issueRequest: Will accept transaction for address 0x" 
                  << std::setfill('0') << std::setw(8) 
                  << std::hex << addr << std::endl;
    } 

    
    // std::cout << "Clock cycle: " << controller->GetClockCycle() << std::endl;
    // controller->AddTransaction(reqId, addr, isWrite);
    uint64_t test_delay = controller->Process(addr, isWrite, reqId, controller->GetClockCycle());
    

    // //参数需要把时间都去掉
    std::cout << "AccurateDRAM issueRequest test_delay: " << test_delay << std::endl;


    // SimpleDRAM原有的代码开始
    int bank = (addr >> lineOffset) & bankMask;
    int row = addr >> rowOffset;

#ifdef __SST_DEBUG_OUTPUT__
    output->debug(_L10_, "AccurateDRAM (%s) received request for address %" PRIx64 " which maps to bank: %d, row: %d. Bank status: %s, open row is %d\n",
           getName().c_str(), addr, bank, row, (busy[bank] ? "busy" : "idle"), openRow[bank]);
#endif

    // if (busy[bank]) return false;

    // int latency = tCAS;
    // if (openRow[bank] != row) {
    //     latency += tRCD;
    //     if (openRow[bank] != -1) {
    //         latency += tRP;
    //         statRowMissRP->addData(1);
    //     } else {
    //         statRowMissNoRP->addData(1);
    //     }
    //     openRow[bank] = row;
    // } else {
    //     statRowHit->addData(1);
    // }
    // busy[bank] = true;
    // std::cout << "SimpleDRAM latency: " << latency << std::endl;
    
    
    // self_link->send(latency, new MemCtrlEvent( bank, reqId ));

    // SimpleDRAM原有的代码结束
    std::cout << "AccurateDRAM issueRequest: Sending MemCtrlEvent with reqId: " << reqId << std::endl;
    self_link->send(test_delay, new MemCtrlEvent( reqId ));

    return true;
}

bool AccurateDRAM::clock(Cycle_t cycle) {
    m_cycle = cycle;
    this->controller->Clock(cycle);
    return false;
}